package proxy

import (
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync/atomic"
	"time"

	"vodka"
	"vodka/logger"
	"vodka/skipper"
)

// TODO: Handle TLS proxy

type (
	// ProxyConfig defines the config for Proxy middleware.
	ProxyConfig struct {
		// Skipper defines a function to skip middleware.
		Skipper skipper.Skipper

		// Balancer defines a load balancing technique.
		// Required.
		// Possible values:
		// - RandomBalancer
		// - RoundRobinBalancer
		Balancer ProxyBalancer
	}

	// ProxyTarget defines the upstream target.
	ProxyTarget struct {
		URL *url.URL
	}

	// RandomBalancer implements a random load balancing technique.
	RandomBalancer struct {
		Targets []*ProxyTarget
		random  *rand.Rand
	}

	// RoundRobinBalancer implements a round-robin load balancing technique.
	RoundRobinBalancer struct {
		Targets []*ProxyTarget
		i       uint32
	}

	// ProxyBalancer defines an interface to implement a load balancing technique.
	ProxyBalancer interface {
		Next() *ProxyTarget
	}
)

func proxyHTTP(t *ProxyTarget) http.Handler {
	return httputil.NewSingleHostReverseProxy(t.URL)
}

/*
type proxyRawHandler struct {
	t *ProxyTarget
	c *vodka.Context
}

func (p *proxyRawHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	in, _, err := p.c.Response.Hijack()
	if err != nil {
		p.c.HandleError(fmt.Errorf("proxy raw, hijack error=%v, url=%s", err, p.t.URL))
		return
	}
	defer in.Close()

	out, err := net.Dial("tcp", p.t.URL.Host)
	if err != nil {
		he := vodka.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, dial error=%v, url=%s", err, p.t.URL))
		p.c.HandleError(he)
		return
	}
	defer out.Close()

	// Write header
	err = r.Write(out)
	if err != nil {
		he := vodka.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, request header copy error=%v, url=%s", p.t.URL, err))
		p.c.HandleError(he)
		return
	}

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	go cp(out, in)
	go cp(in, out)
	err = <-errc
	if err != nil && err != io.EOF {
		fmt.Errorf("proxy raw, copy body error=%v, url=%s", err, p.t.URL)
	}
}
*/

func proxyRaw(t *ProxyTarget, c *vodka.Context) http.Handler {
	//return &proxyRawHandler{t, c}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		in, _, err := c.Response.Hijack()
		if err != nil {
			c.HandleError(fmt.Errorf("proxy raw, hijack url=%s, error=%v", t.URL, err))
			return
		}
		defer in.Close()

		out, err := net.Dial("tcp", t.URL.Host)
		if err != nil {
			he := vodka.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, dial url=%s, error=%v", t.URL, err))
			c.HandleError(he)
			return
		}
		defer out.Close()

		// Write header
		err = r.Write(out)
		if err != nil {
			he := vodka.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, request header copy url=%s, error=%v", t.URL, err))
			c.HandleError(he)
			return
		}

		errc := make(chan error, 2)
		cp := func(dst io.Writer, src io.Reader) {
			_, err := io.Copy(dst, src)
			errc <- err
		}

		go cp(out, in)
		go cp(in, out)
		err = <-errc
		if err != nil && err != io.EOF {
			fmt.Errorf("proxy raw, copy body url=%s, error=%v", t.URL, err)
		}
	})
}

// Next randomly returns an upstream target.
func (r *RandomBalancer) Next() *ProxyTarget {
	if r.random == nil {
		r.random = rand.New(rand.NewSource(int64(time.Now().Nanosecond())))
	}
	return r.Targets[r.random.Intn(len(r.Targets))]
}

// Next returns an upstream target using round-robin technique.
func (r *RoundRobinBalancer) Next() *ProxyTarget {
	r.i = r.i % uint32(len(r.Targets))
	t := r.Targets[r.i]
	atomic.AddUint32(&r.i, 1)
	return t
}

// Proxy returns an HTTP/WebSocket reverse proxy middleware.
func Proxy(config ProxyConfig) vodka.Handler {
	// Defaults
	if config.Skipper == nil {
		config.Skipper = logger.DefaultLoggerConfig.Skipper
	}
	if config.Balancer == nil {
		panic("vodka: proxy middleware requires balancer")
	}

	return func(c *vodka.Context) (err error) {
		req := c.Request
		res := c.Response
		tgt := config.Balancer.Next()

		// Fix header
		if req.Header.Get(vodka.HeaderXRealIP) == "" {
			req.Header.Set(vodka.HeaderXRealIP, c.RealIP())
		}
		if req.Header.Get(vodka.HeaderXForwardedProto) == "" {
			req.Header.Set(vodka.HeaderXForwardedProto, c.Scheme())
		}
		if c.IsWebSocket() && req.Header.Get(vodka.HeaderXForwardedFor) == "" { // For HTTP, it is automatically set by Go HTTP reverse proxy.
			req.Header.Set(vodka.HeaderXForwardedFor, c.RealIP())
		}

		// Proxy
		switch {
		case c.IsWebSocket():
			proxyRaw(tgt, c).ServeHTTP(res, req)
		case req.Header.Get(vodka.HeaderAccept) == "text/event-stream":
		default:
			proxyHTTP(tgt).ServeHTTP(res, req)
		}

		return c.Abort()
	}

}
